import os
import os.path as osp
from pathlib import Path
import multiprocessing as mp
from multiprocessing.pool import Pool
import json
from copy import deepcopy
from collections import defaultdict

import numpy as np
import cv2 as cv


IMAGE_SUFFIX = ['.jpg', '.jpeg', '.png']

COLOR_MAP = {
    "background": (0, 0, 0),
    "road_surface": (240, 240, 70),
    "parking_slot_line": (0, 0, 128),
    "lane_line": (255, 190, 220),
    "guide_arrow": (0, 128, 128),
    "zebra_crossing": (75, 25, 230),
    "speed_bump": (230, 50, 240),
    "parking_slot_id": (128, 128, 0),
    "road_text": (180, 215, 255),
    "wheel_blocker": (25, 255, 255),
    "other_roader_marker": (195, 255, 170),
    "wall": (128, 0, 0),
    "pillar": (48, 130, 245),
    "ego_car": (108, 69, 166),
    "clean_no_parking_area": (75, 180, 60),
    "curb": (80, 80, 80),
    "mesh_no_stop_zone": (100, 100, 100),
    "stop_line": (120, 120, 120),
    "vehicle": (150, 150, 150),
    }


def read_json(path, mode='r'):
    with open(path, mode, encoding='utf8') as fp:
        data = json.load(fp)
    return data

def cv_show(window_name, image_data):
    cv.namedWindow(window_name, cv.WINDOW_NORMAL)
    cv.imshow(window_name, image_data)
    key = cv.waitKey(0)
    if key == "q":
        cv.destroyAllWindows()
    if key == 27:
        cv.destroyAllWindows()
        exit(0)


def obt_points(label_data):
    # name_list = list()
    # allpoints_list = list()
    point_dict = dict()
    for obj in label_data["LabelDetail"]["objects"]:
        points_list = list()
        for point in obj["feature"]["points"]:
            points_list.append([int(float(point["x"])), int(float(point["y"]))])
        if not obj["labelCode"] in point_dict:
            point_dict[obj["labelCode"]] = list()
            point_dict[obj["labelCode"]].append(np.array(points_list))
        else:
            point_dict[obj["labelCode"]].append(np.array(points_list))
        
        
        # if obj["labelCode"] == "road_surface":
        #     name_list.insert(0, obj["labelCode"])
        #     allpoints_list.insert(0, np.array(points_list))
        # else:
        #     name_list.append(obj["labelCode"])
        #     allpoints_list.append(np.array(points_list))
    return point_dict

def rander(image, points):
    # for name, point in points.items():
    #     # print(f"===> points_name: {name}  |*-*|  points: {points}")
    #     cv.fillPoly(image, [point], COLOR_MAP[name])
    #     # cv.rectangle(image, (np.min(point, axis=0)), (np.max(point, axis=0)), COLOR_MAP[name], 2)
    # # print("\n")
    for name in COLOR_MAP:
        if name in points:
            for point in points[name]:
                cv.fillPoly(image, [point], COLOR_MAP[name])
                cv.rectangle(image, np.min(point, axis=0), np.max(point, axis=0), COLOR_MAP[name], 2)
    return image

def process(label_path, image_path, save_path):
    image = cv.imread(image_path)
    label_data = read_json(label_path)
    points = obt_points(label_data)
    result = rander(image, points)
    print(f"===> save mask to: '{save_path}'")
    cv.imwrite(save_path, result)


if __name__ == "__main__":
    image_dir = "img"
    label_dir = "json"
    save_dir = "mask_3"
    
    label_path_list = sorted([path for path in Path(label_dir).rglob("*") if path.is_file() and path.suffix == ".json"])
    filename_list = sorted([path.stem for path in label_path_list])
    image_path_list = sorted([str(path) for path in Path(image_dir).rglob("*") if path.is_file() and path.suffix in IMAGE_SUFFIX and path.stem in filename_list])
    save_path_list = []
    for path in image_path_list:
        save_path = path.replace(image_dir, save_dir)
        if not osp.exists(osp.dirname(save_path)):
            os.makedirs(osp.dirname(save_path))
        save_path_list.append(save_path)
        save_path_list = sorted(save_path_list)
    
    print(f"""
        {'*' * 60}
            Load images from: '{osp.abspath(image_dir)}'
            Load labels from: '{osp.abspath(label_dir)}'
            Save masks to: {osp.abspath(save_dir)}
            All process total numbers: {len(label_path_list)}
        {'*' * 60}
            """)
    
    # # 多进程处理
    # with Pool(mp.cpu_count()) as rander_pool:
    #     rander_pool.starmap(process, zip(label_path_list, image_path_list, save_path_list))
    
    # 调试 可视化
    for label_path, image_path in zip(label_path_list, image_path_list):
        # if osp.basename(label_path) in ["20230716113738.930964_SurCam02.json", "20230716113738.930964_SurCam01.json"]: 
        image = cv.imread(image_path)
        image_copy = deepcopy(image)
        label_data = read_json(label_path)
        points = obt_points(label_data)
        result = rander(image_copy, points)
        image_stack = np.hstack([image, result])
        cv_show("stack", image_stack)
